Skip to content

Update Jax code#244

Merged
ludgerpaehler merged 14 commits intomainfrom
cl/JAX_update
Apr 24, 2026
Merged

Update Jax code#244
ludgerpaehler merged 14 commits intomainfrom
cl/JAX_update

Conversation

@cl126162
Copy link
Copy Markdown
Collaborator

@cl126162 cl126162 commented Apr 7, 2026

Kolmogorov flow:

  • refactor obs extracttion avoiding multiple FFT computations
  • add full jit support through entire env
  • FP32 support: stable at dt=7.5e-4 | FP64 support: stable at dt=1e-3

Channel flow:

  • full fp32 support
  • add full jit support through entire env

@smokbel Can you please carefully double-check if physical correctness is preserved?

@ludgerpaehler
Copy link
Copy Markdown
Collaborator

Any chance we can at least separate out the Nek components from this PR? Also, can we try to move the .yml configurations out into the HuggingFace repository?

@cl126162
Copy link
Copy Markdown
Collaborator Author

Configs are pushed now to HF

@cl126162
Copy link
Copy Markdown
Collaborator Author

nek and maia updates are now separated

@ludgerpaehler ludgerpaehler merged commit 422799a into main Apr 24, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants